import cv2
import numpy as np
import os

def calculate_straightness(points):
    # 第一个角点和最后一个角点
    point1 = points[0]
    point2 = points[-1]
    
    # 计算这两个点之间的距离
    dist = np.linalg.norm(point2 - point1)
    
    # 计算以这两个点为边长的等边三角形的第三点
    midpoint = (point1 + point2) / 2
    height = dist * np.sin(np.radians(60))
    direction = (point2 - point1) / dist
    perpendicular_direction = np.array([-direction[1], direction[0]])
    third_point = midpoint + height * perpendicular_direction
    
    # 构建等边三角形的顶点
    equilateral_triangle = np.array([point1, point2, third_point])
    
    # 计算等边三角形的面积
    area_equilateral = cv2.contourArea(equilateral_triangle)
    
    # 计算前num_points个点形成的多边形的面积
    polygon_area = cv2.contourArea(points)
    
    # 计算直线度
    straightness = (area_equilateral - polygon_area) / area_equilateral
    
    return straightness

def detect_and_annotate_chessboard(image_path, output_dir, pattern_size=(6, 9)):
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 读取图像
    image = cv2.imread(image_path)
    if image is None:
        raise FileNotFoundError(f"Image not found at {image_path}")
    
    # 保存原始图像
    cv2.imwrite(os.path.join(output_dir, "original.jpg"), image)
    print(f"Saved original image as {os.path.join(output_dir, 'original.jpg')}")
    
    # 转换为灰度图像
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    
    # 保存灰度图像
    cv2.imwrite(os.path.join(output_dir, "gray.jpg"), gray)
    print(f"Saved gray image as {os.path.join(output_dir, 'gray.jpg')}")
    
    # 设置角点检测的迭代终止条件
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    
    # 检测棋盘格角点
    ret, corners = cv2.findChessboardCorners(gray, pattern_size, None)
    # 检测直线度
    print(calculate_straightness(corners, 6))
    # 如果检测到角点
    if ret:
        print("Chessboard corners detected.")
        
        
        
        # 细化角点
        corners_subpix = cv2.cornerSubPix(gray, corners, (12, 8), (-1, -1), criteria)
        
        # 保存细化前的角点图像（用红色圆点标记）
        image_with_raw_corners = cv2.drawChessboardCorners(image.copy(), pattern_size, corners, ret)
        cv2.imwrite(os.path.join(output_dir, "raw_corners.jpg"), image_with_raw_corners)
        print(f"Saved raw corners image as {os.path.join(output_dir, 'raw_corners.jpg')}")
        
        # 保存细化后的角点图像（用绿色圆点标记）
        image_with_subpix_corners = cv2.drawChessboardCorners(image.copy(), pattern_size, corners_subpix, ret)
        cv2.imwrite(os.path.join(output_dir, "subpix_corners.jpg"), image_with_subpix_corners)
        print(f"Saved subpixel corners image as {os.path.join(output_dir, 'subpix_corners.jpg')}")
        
        # 输出角点坐标
        print(f"Chessboard corners coordinates (in pixels):")
        # for corner in corners_subpix.astype(int):  # 转换为整数坐标以便打印
        #     print(corner.ravel())  # 将二维坐标转换为一维数组并打印
    else:
        print("Chessboard corners not found.")

# 调用函数，传入图像路径和输出目录
image_path = 'images_processing/img_out/undistorted_img.jpg'  # 替换为你的棋盘格图像路径
output_dir = 'images_processing/img_out/corner/'  # 替换为你希望保存输出图像的目录
detect_and_annotate_chessboard(image_path, output_dir)